#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <thread>
#include <mutex>
#include <chrono>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <ctime>
#include <cstdlib>
#include <random>
#include "sqlite3.h"

#pragma comment(lib, "ws2_32.lib")
// #pragma comment(lib, "sqlite3.lib")

#define PORT 8080
#define HEARTBEAT_INTERVAL 10 // 30分钟(秒)
#define DB_FILE "licenses.db"

sqlite3 *db;
std::mutex dbMutex;

bool initializeDatabase()
{
    int rc = sqlite3_open(DB_FILE, &db);
    if (rc != SQLITE_OK)
    {
        std::cerr << "无法打开数据库: " << sqlite3_errmsg(db) << std::endl;
        return false;
    }

    const char *createLicenseTable =
        "CREATE TABLE IF NOT EXISTS licenses ("
        "serial TEXT PRIMARY KEY,"
        "max_users INTEGER NOT NULL,"
        "active_users INTEGER NOT NULL DEFAULT 0);";

    const char *createActiveUserTable =
        "CREATE TABLE IF NOT EXISTS active_users ("
        "serial TEXT,"
        "userid TEXT,"
        "last_check INTEGER NOT NULL,"
        "PRIMARY KEY (serial, userid),"
        "FOREIGN KEY (serial) REFERENCES licenses(serial));";

    char *errMsg = nullptr;
    rc = sqlite3_exec(db, createLicenseTable, nullptr, nullptr, &errMsg);
    if (rc != SQLITE_OK)
    {
        std::cerr << "创建许可证表失败: " << errMsg << std::endl;
        sqlite3_free(errMsg);
        return false;
    }

    rc = sqlite3_exec(db, createActiveUserTable, nullptr, nullptr, &errMsg);
    if (rc != SQLITE_OK)
    {
        std::cerr << "创建活跃用户表失败: " << errMsg << std::endl;
        sqlite3_free(errMsg);
        return false;
    }

    return true;
}

std::string generateRandomString(int length, const std::string &charset = "0123456789")
{
    static std::random_device rd;
    static std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, charset.size() - 1);

    std::string result;
    for (int i = 0; i < length; i++)
    {
        result += charset[dis(gen)];
    }
    return result;
}

std::string initializeLicense(const std::string &username, int maxUsers)
{
    std::lock_guard<std::mutex> lock(dbMutex);

    std::string serial;
    sqlite3_stmt *stmt;

    do
    {
        serial = generateRandomString(10);
        const char *checkSql = "SELECT serial FROM licenses WHERE serial = ?;";
        if (sqlite3_prepare_v2(db, checkSql, -1, &stmt, nullptr) != SQLITE_OK)
        {
            std::cerr << "准备SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
            return "";
        }

        sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
        int rc = sqlite3_step(stmt);
        sqlite3_finalize(stmt);

        if (rc == SQLITE_DONE)
            break;
    } while (true);

    const char *insertSql = "INSERT INTO licenses (serial, max_users, active_users) VALUES (?, ?, 0);";
    if (sqlite3_prepare_v2(db, insertSql, -1, &stmt, nullptr) != SQLITE_OK)
    {
        std::cerr << "准备插入SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
        return "";
    }

    sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_int(stmt, 2, maxUsers);

    if (sqlite3_step(stmt) != SQLITE_DONE)
    {
        std::cerr << "插入许可证失败: " << sqlite3_errmsg(db) << std::endl;
        sqlite3_finalize(stmt);
        return "";
    }

    sqlite3_finalize(stmt);
    std::cout << "License generated for " << username << ": " << serial << " (Max Users: " << maxUsers << ")" << std::endl;
    return serial;
}

std::string validateLicense(const std::string &serial)
{
    std::lock_guard<std::mutex> lock(dbMutex);
    sqlite3_stmt *stmt;

    const char *checkSql = "SELECT max_users, active_users FROM licenses WHERE serial = ?;";
    if (sqlite3_prepare_v2(db, checkSql, -1, &stmt, nullptr) != SQLITE_OK)
    {
        std::cerr << "准备验证SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
        return "DATABASE ERROR";
    }

    sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);

    if (sqlite3_step(stmt) != SQLITE_ROW)
    {
        sqlite3_finalize(stmt);
        return "WRONG SERIAL!";
    }

    int maxUsers = sqlite3_column_int(stmt, 0);
    int activeUsers = sqlite3_column_int(stmt, 1);
    sqlite3_finalize(stmt);

    if (activeUsers >= maxUsers)
    {
        return "MAX users reached. WAIT or contact admin for UPGRADE!";
    }

    std::string userid;
    do
    {
        userid = generateRandomString(6);
        const char *checkUserSql = "SELECT userid FROM active_users WHERE serial = ? AND userid = ?;";
        if (sqlite3_prepare_v2(db, checkUserSql, -1, &stmt, nullptr) != SQLITE_OK)
        {
            std::cerr << "准备用户检查SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
            return "DATABASE ERROR";
        }

        sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
        sqlite3_bind_text(stmt, 2, userid.c_str(), -1, SQLITE_STATIC);

        int rc = sqlite3_step(stmt);
        sqlite3_finalize(stmt);

        if (rc == SQLITE_DONE)
            break;
    } while (true);

    const char *updateSql = "UPDATE licenses SET active_users = active_users + 1 WHERE serial = ?;";
    if (sqlite3_prepare_v2(db, updateSql, -1, &stmt, nullptr) != SQLITE_OK)
    {
        std::cerr << "准备更新SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
        return "DATABASE ERROR";
    }

    sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
    if (sqlite3_step(stmt) != SQLITE_DONE)
    {
        sqlite3_finalize(stmt);
        return "DATABASE ERROR";
    }
    sqlite3_finalize(stmt);

    const char *insertUserSql = "INSERT INTO active_users (serial, userid, last_check) VALUES (?, ?, ?);";
    if (sqlite3_prepare_v2(db, insertUserSql, -1, &stmt, nullptr) != SQLITE_OK)
    {
        std::cerr << "准备插入用户SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
        return "DATABASE ERROR";
    }

    std::time_t now = std::time(nullptr);
    sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 2, userid.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_int64(stmt, 3, now);

    if (sqlite3_step(stmt) != SQLITE_DONE)
    {
        sqlite3_finalize(stmt);
        return "DATABASE ERROR";
    }

    sqlite3_finalize(stmt);
    std::cerr << "usertotal:" << activeUsers + 1 << std::endl;
    return userid;
}

void updateHeartbeat(const std::string &serial, const std::string &userid)
{
    std::lock_guard<std::mutex> lock(dbMutex);
    sqlite3_stmt *stmt;

    const char *updateSql = "UPDATE active_users SET last_check = ? WHERE serial = ? AND userid = ?;";
    if (sqlite3_prepare_v2(db, updateSql, -1, &stmt, nullptr) != SQLITE_OK)
    {
        std::cerr << "准备心跳更新SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
        return;
    }

    std::time_t now = std::time(nullptr);
    sqlite3_bind_int64(stmt, 1, now);
    sqlite3_bind_text(stmt, 2, serial.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 3, userid.c_str(), -1, SQLITE_STATIC);

    if (sqlite3_step(stmt) != SQLITE_DONE)
    {
        std::cerr << "更新心跳失败: " << sqlite3_errmsg(db) << std::endl;
    }

    sqlite3_finalize(stmt);
}

void checkHeartbeats()
{
    while (true)
    {
        std::this_thread::sleep_for(std::chrono::seconds(HEARTBEAT_INTERVAL));
        std::lock_guard<std::mutex> lock(dbMutex);

        std::time_t now = std::time(nullptr);
        std::time_t threshold = now - HEARTBEAT_INTERVAL;

        const char *selectTimeoutSql =
            "SELECT serial, userid FROM active_users WHERE last_check < ?;";

        sqlite3_stmt *stmt;
        if (sqlite3_prepare_v2(db, selectTimeoutSql, -1, &stmt, nullptr) != SQLITE_OK)
        {
            std::cerr << "准备超时查询SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
            continue;
        }

        sqlite3_bind_int64(stmt, 1, threshold);

        std::vector<std::pair<std::string, std::string>> timeoutUsers;
        while (sqlite3_step(stmt) == SQLITE_ROW)
        {
            const char *serial = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 0));
            const char *userid = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 1));
            timeoutUsers.emplace_back(serial, userid);
        }
        sqlite3_finalize(stmt);

        for (const auto &[serial, userid] : timeoutUsers)
        {
            const char *deleteUserSql = "DELETE FROM active_users WHERE serial = ? AND userid = ?;";
            if (sqlite3_prepare_v2(db, deleteUserSql, -1, &stmt, nullptr) != SQLITE_OK)
            {
                std::cerr << "准备删除用户SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
                continue;
            }

            sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_text(stmt, 2, userid.c_str(), -1, SQLITE_STATIC);

            if (sqlite3_step(stmt) != SQLITE_DONE)
            {
                std::cerr << "删除用户失败: " << sqlite3_errmsg(db) << std::endl;
                sqlite3_finalize(stmt);
                continue;
            }
            sqlite3_finalize(stmt);

            const char *updateCountSql = "UPDATE licenses SET active_users = active_users - 1 WHERE serial = ?;";
            if (sqlite3_prepare_v2(db, updateCountSql, -1, &stmt, nullptr) != SQLITE_OK)
            {
                std::cerr << "准备更新计数SQL语句失败: " << sqlite3_errmsg(db) << std::endl;
                continue;
            }

            sqlite3_bind_text(stmt, 1, serial.c_str(), -1, SQLITE_STATIC);
            if (sqlite3_step(stmt) != SQLITE_DONE)
            {
                std::cerr << "更新计数失败: " << sqlite3_errmsg(db) << std::endl;
            }
            sqlite3_finalize(stmt);

            std::cout << "License: " << serial << " userid: " << userid << ": removed" << std::endl;
        }
    }
}

void handleClient(SOCKET clientSocket)
{
    char buffer[1024] = {0};
    while (true)
    {
        int bytesRead = recv(clientSocket, buffer, sizeof(buffer) - 1, 0);
        if (bytesRead <= 0)
        {
            break;
        }

        buffer[bytesRead] = '\0';
        std::string request(buffer);
        std::string response;

        if (request.substr(0, 9) == "VALIDATE:")
        {
            std::string serial = request.substr(9);
            std::cerr << serial << std::endl;
            std::string userid = validateLicense(serial);
            if (userid.size() == 6)
            {
                response = "AUTHORIZED:" + userid;
            }
            else
            {
                response = "DENIED:" + userid;
            }
        }
        else if (request.substr(0, 10) == "HEARTBEAT:")
        {
            std::string serial = request.substr(10, 10);
            std::string userid = request.substr(21, 6);
            updateHeartbeat(serial, userid);
            response = "HEARTBEAT_ACK";
        }
        else if (request.substr(0, 9) == "PURCHASE:")
        {
            size_t pos1 = request.find(':', 8);
            size_t pos2 = request.find(':', pos1 + 1);
            size_t pos3 = request.find(':', pos2 + 1);

            if (pos1 != std::string::npos && pos2 != std::string::npos && pos3 != std::string::npos)
            {
                std::string username = request.substr(pos1 + 1, pos2 - pos1 - 1);
                std::string key = request.substr(pos2 + 1, pos3 - pos2 - 1);
                std::string version = request.substr(pos3 + 1);

                if (username == "admin" && key == "123456")
                {
                    std::string password = initializeLicense(username, version == "A" ? 50 : 10);
                    response = "SUCCESS:" + password;
                }
                else
                {
                    response = "WRONG USERNAME OR KEY !";
                }
            }
        }

        send(clientSocket, response.c_str(), response.size(), 0);
    }
    closesocket(clientSocket);
}

int main()
{
    if (!initializeDatabase())
    {
        std::cerr << "数据库初始化失败" << std::endl;
        return -1;
    }

    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
        return -1;
    }

    std::thread heartbeatThread(checkHeartbeats);
    heartbeatThread.detach();

    SOCKET serverSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (serverSocket == INVALID_SOCKET)
    {
        std::cerr << "Socket creation failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return -1;
    }

    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = INADDR_ANY;
    serverAddr.sin_port = htons(PORT);

    if (bind(serverSocket, (sockaddr *)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR)
    {
        std::cerr << "Bind failed: " << WSAGetLastError() << std::endl;
        closesocket(serverSocket);
        WSACleanup();
        return -1;
    }

    if (listen(serverSocket, SOMAXCONN) == SOCKET_ERROR)
    {
        std::cerr << "Listen failed: " << WSAGetLastError() << std::endl;
        closesocket(serverSocket);
        WSACleanup();
        return -1;
    }

    std::cout << "Server started on port " << PORT << std::endl;

    while (true)
    {
        SOCKET clientSocket = accept(serverSocket, nullptr, nullptr);
        if (clientSocket == INVALID_SOCKET)
        {
            std::cerr << "Accept failed: " << WSAGetLastError() << std::endl;
            continue;
        }

        std::thread clientThread(handleClient, clientSocket);
        clientThread.detach();
    }

    closesocket(serverSocket);
    sqlite3_close(db);
    WSACleanup();
    return 0;
}